{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "beObUOFyuRjT"
      },
      "source": [
        "##### Copyright 2018 The TF-Agents Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "nQnmcm0oI1Q-"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eutDVTs9aJEL"
      },
      "source": [
        "# 再生バッファ\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://www.tensorflow.org/agents/tutorials/5_replay_buffers_tutorial\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\"> TensorFlow.orgで表示</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ja/agents/tutorials/5_replay_buffers_tutorial.ipynb\">     <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Google Colabで実行</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ja/agents/tutorials/5_replay_buffers_tutorial.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">GitHub でソースを表示{</a></td>\n",
        "  <td><a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ja/agents/tutorials/5_replay_buffers_tutorial.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\">ノートブックをダウンロード/a0}</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8aPHF9kXFggA"
      },
      "source": [
        "## はじめに\n",
        "\n",
        "強化学習アルゴリズムは、環境でポリシーを実行するときに再生バッファを使用して経験のトラジェクトリを格納します。トレーニング中、再生バッファはトラジェクトリのサブセット（シーケンシャルサブセットまたはサンプルのいずれか）がエージェントの経験を「再生」するために照会されます。\n",
        "\n",
        "このコラボでは、共通の API を共有する 2 種類の再生バッファ（python 方式と tensorflow 方式）について考察します。次のセクションでは API と各バッファの実装、およびデータ収集トレーニングでそれらを使用する方法を説明します。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1uSlqYgvaG9b"
      },
      "source": [
        "## セットアップ"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GztmUpWKZ7kq"
      },
      "source": [
        "tf-agents をまだインストールしていない場合はインストールしてください。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TnE2CgilrngG"
      },
      "outputs": [],
      "source": [
        "!pip install tf-agents\n",
        "!pip install gym"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "whYNP894FSkA"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import\n",
        "from __future__ import division\n",
        "from __future__ import print_function\n",
        "\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "\n",
        "from tf_agents import specs\n",
        "from tf_agents.agents.dqn import dqn_agent\n",
        "from tf_agents.drivers import dynamic_step_driver\n",
        "from tf_agents.environments import suite_gym\n",
        "from tf_agents.environments import tf_py_environment\n",
        "from tf_agents.networks import q_network\n",
        "from tf_agents.replay_buffers import py_uniform_replay_buffer\n",
        "from tf_agents.replay_buffers import tf_uniform_replay_buffer\n",
        "from tf_agents.specs import tensor_spec\n",
        "from tf_agents.trajectories import time_step\n",
        "\n",
        "tf.compat.v1.enable_v2_behavior()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xcQWclL9FpZl"
      },
      "source": [
        "## 再生バッファ API\n",
        "\n",
        "再生バッファのクラスには、次の定義とメソッドがあります。\n",
        "\n",
        "```python\n",
        "class ReplayBuffer(tf.Module):   \"\"\"Abstract base class for TF-Agents replay buffer.\"\"\"    def __init__(self, data_spec, capacity):     \"\"\"Initializes the replay buffer.      Args:       data_spec: A spec or a list/tuple/nest of specs describing         a single item that can be stored in this buffer       capacity: number of elements that the replay buffer can hold.     \"\"\"    @property   def data_spec(self):     \"\"\"Returns the spec for items in the replay buffer.\"\"\"    @property   def capacity(self):     \"\"\"Returns the capacity of the replay buffer.\"\"\"    def add_batch(self, items):     \"\"\"Adds a batch of items to the replay buffer.\"\"\"    def get_next(self,                sample_batch_size=None,                num_steps=None,                time_stacked=True):     \"\"\"Returns an item or batch of items from the buffer.\"\"\"    def as_dataset(self,                  sample_batch_size=None,                  num_steps=None,                  num_parallel_calls=None):     \"\"\"Creates and returns a dataset that returns entries from the buffer.\"\"\"     def gather_all(self):     \"\"\"Returns all the items in buffer.\"\"\"     return self._gather_all()    def clear(self):     \"\"\"Resets the contents of replay buffer\"\"\"\n",
        "```\n",
        "\n",
        "再生オブジェクトバッファのオブジェクトが初期化される際には、そのオブジェクトが格納する要素の `data_spec` が必要です。この仕様は、バッファに追加されるトラジェクトリ要素の `TensorSpec` に対応しています。この仕様は通常、トレーニング時にエージェントが期待する形状、タイプ、構造を定義するエージェントの `agent.collect_data_spec` を調査することで取得できます（詳細は後述します）。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X3Yrxg36Ik1x"
      },
      "source": [
        "## TFUniformReplayBuffer\n",
        "\n",
        "`TFUniformReplayBuffer` は TF-Agents で最も一般的に使用される再生バッファであるため、このチュートリアルではこの再生バッファを使用します。`TFUniformReplayBuffer` では、バッキングバッファの記憶は tensorflow 変数によって行われるため、計算グラフの一部になっています。\n",
        "\n",
        "バッファには複数の要素がバッチ単位で格納され、バッチセグメントごとに最大容量の `max_length` 要素があります。したがって、合計バッファ容量は、`batch_size` x `max_length` 要素となります。バッファに格納される要素はすべて、対応するデータ仕様を持っている必要があります。再生バッファがデータ収集に使用される場合、仕様はエージェントの収集データ仕様になります。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lYk-bn2taXlw"
      },
      "source": [
        "### バッファの作成:\n",
        "\n",
        "`TFUniformReplayBuffer` を作成するには、次の値を渡します。\n",
        "\n",
        "1. バッファが格納するデータ要素の仕様\n",
        "2. バッファのバッチサイズに対応する `batch size`\n",
        "3. バッチセグメントごとの `max_length` 要素数\n",
        "\n",
        "サンプルデータ仕様（`batch_size` 32 および `max_length` 1000）を使用して `TFUniformReplayBuffer` を作成する例を以下に示します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dj4_-77_5ExP"
      },
      "outputs": [],
      "source": [
        "data_spec =  (\n",
        "        tf.TensorSpec([3], tf.float32, 'action'),\n",
        "        (\n",
        "            tf.TensorSpec([5], tf.float32, 'lidar'),\n",
        "            tf.TensorSpec([3, 2], tf.float32, 'camera')\n",
        "        )\n",
        ")\n",
        "\n",
        "batch_size = 32\n",
        "max_length = 1000\n",
        "\n",
        "replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(\n",
        "    data_spec,\n",
        "    batch_size=batch_size,\n",
        "    max_length=max_length)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XB8rOw5ATDD2"
      },
      "source": [
        "### バッファへの書き込み:\n",
        "\n",
        "再生バッファに要素を追加するため、`items` がバッファに追加されるアイテムのバッチを表すテンソルのリスト/タプル/ネストになっている `add_batch(items)` メソッドを使用します。`items` の各要素には外側の次元に等しい `batch_size` が必要であり、残りの次元は（再生バッファのコンストラクタに渡されたデータ仕様と同じ）アイテムのデータ仕様に準拠していなければなりません。\n",
        "\n",
        "アイテムのバッチを追加する例を以下に示します。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nOvkp4vJhBOT"
      },
      "outputs": [],
      "source": [
        "action = tf.constant(1 * np.ones(\n",
        "    data_spec[0].shape.as_list(), dtype=np.float32))\n",
        "lidar = tf.constant(\n",
        "    2 * np.ones(data_spec[1][0].shape.as_list(), dtype=np.float32))\n",
        "camera = tf.constant(\n",
        "    3 * np.ones(data_spec[1][1].shape.as_list(), dtype=np.float32))\n",
        "  \n",
        "values = (action, (lidar, camera))\n",
        "values_batched = tf.nest.map_structure(lambda t: tf.stack([t] * batch_size),\n",
        "                                       values)\n",
        "  \n",
        "replay_buffer.add_batch(values_batched)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "smnVAxHghKly"
      },
      "source": [
        "### バッファからの読み込み\n",
        "\n",
        "`TFUniformReplayBuffer` からデータを読み込む方法は 3 つあります。\n",
        "\n",
        "1. `get_next()` - バッファからサンプルを 1 つ返します。返されるサンプルバッチサイズとタイムステップ数は、このメソッドの引数で指定できます。\n",
        "2. `as_dataset()` - 再生バッファを `tf.data.Dataset` として返します。その後、データセットイテレータを作成し、バッファ内のアイテムのサンプルをイテレートできます。\n",
        "3. `gather_all()` - バッファ内のすべてのアイテムを形状 `[batch, time, data_spec]` を含むテンソルとして返します。\n",
        "\n",
        "上記の各メソッドを使用して再生バッファから読み込む例を以下に示します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IlQ1eGhohM3M"
      },
      "outputs": [],
      "source": [
        "# add more items to the buffer before reading\n",
        "for _ in range(5):\n",
        "  replay_buffer.add_batch(values_batched)\n",
        "\n",
        "# Get one sample from the replay buffer with batch size 10 and 1 timestep:\n",
        "\n",
        "sample = replay_buffer.get_next(sample_batch_size=10, num_steps=1)\n",
        "\n",
        "# Convert the replay buffer to a tf.data.Dataset and iterate through it\n",
        "dataset = replay_buffer.as_dataset(\n",
        "    sample_batch_size=4,\n",
        "    num_steps=2)\n",
        "\n",
        "iterator = iter(dataset)\n",
        "print(\"Iterator trajectories:\")\n",
        "trajectories = []\n",
        "for _ in range(3):\n",
        "  t, _ = next(iterator)\n",
        "  trajectories.append(t)\n",
        "  \n",
        "print(tf.nest.map_structure(lambda t: t.shape, trajectories))\n",
        "\n",
        "# Read all elements in the replay buffer:\n",
        "trajectories = replay_buffer.gather_all()\n",
        "\n",
        "print(\"Trajectories from gather all:\")\n",
        "print(tf.nest.map_structure(lambda t: t.shape, trajectories))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BcS49HrNF34W"
      },
      "source": [
        "## PyUniformReplayBuffer\n",
        "\n",
        "`PyUniformReplayBuffer` は `TFUniformReplayBuffer` と同じ機能を持ちますが、データは tf 変数ではなく numpy 配列に格納されます。このバッファはグラフ外のデータ収集に使用されます。numpy にバッキングストレージを用意すると、Tensorflow 変数を使用せずに一部のアプリケーションがデータ操作（優先度を更新するためのインデックス作成など）を実行しやすくなります。ただし、この実装では Tensorflow を使用したグラフ最適化のメリットを享受できません。\n",
        "\n",
        "エージェントのポリシートラジェクトリの仕様から `PyUniformReplayBuffer` をインスタンス化する例を以下に示します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F4neLPpL25wI"
      },
      "outputs": [],
      "source": [
        "replay_buffer_capacity = 1000*32 # same capacity as the TFUniformReplayBuffer\n",
        "\n",
        "py_replay_buffer = py_uniform_replay_buffer.PyUniformReplayBuffer(\n",
        "    capacity=replay_buffer_capacity,\n",
        "    data_spec=tensor_spec.to_nest_array_spec(data_spec))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9V7DEcB8IeiQ"
      },
      "source": [
        "## トレーニング中に再生バッファを使用する\n",
        "\n",
        "再生バッファを作成してアイテムを読み書きする方法がわかったので、エージェントのトレーニング中に再生バッファを使用してトラジェクトリを格納できるようになりました。\n",
        "\n",
        "### データ収集\n",
        "\n",
        "まず、データ収集中に再生バッファを使用する方法を見てみましょう。\n",
        "\n",
        "TF-Agents では、`Driver`（詳細は Driver のチュートリアルをご覧ください）を使用して環境内の経験を収集します。`Driver` を使用するには、`Driver` がトラジェクトリを受け取ったときに実行する関数である `Observer` を指定します。\n",
        "\n",
        "そのため、トラジェクトリ要素を再生バッファに追加するために `add_batch(items)` を呼び出すオブザーバーを追加してアイテム（のバッチ）を再生バッファに追加します。\n",
        "\n",
        "`TFUniformReplayBuffer` を使用した例を以下に示します。まず、環境、ネットワーク、エージェントを作成します。その後、`TFUniformReplayBuffer` を作成します。再生バッファ内のトラジェクトリ要素の仕様は、エージェントの収集データ仕様に等しくなっていることに注意してください。その後、その `add_batch` メソッドをトレーニング中にデータ収集を実行するドライバーのオブザーバーに設定します。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pCbTDO3Z5UCS"
      },
      "outputs": [],
      "source": [
        "env = suite_gym.load('CartPole-v0')\n",
        "tf_env = tf_py_environment.TFPyEnvironment(env)\n",
        "\n",
        "q_net = q_network.QNetwork(\n",
        "    tf_env.time_step_spec().observation,\n",
        "    tf_env.action_spec(),\n",
        "    fc_layer_params=(100,))\n",
        "\n",
        "agent = dqn_agent.DqnAgent(\n",
        "    tf_env.time_step_spec(),\n",
        "    tf_env.action_spec(),\n",
        "    q_network=q_net,\n",
        "    optimizer=tf.compat.v1.train.AdamOptimizer(0.001))\n",
        "\n",
        "replay_buffer_capacity = 1000\n",
        "\n",
        "replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(\n",
        "    agent.collect_data_spec,\n",
        "    batch_size=tf_env.batch_size,\n",
        "    max_length=replay_buffer_capacity)\n",
        "\n",
        "# Add an observer that adds to the replay buffer:\n",
        "replay_observer = [replay_buffer.add_batch]\n",
        "\n",
        "collect_steps_per_iteration = 10\n",
        "collect_op = dynamic_step_driver.DynamicStepDriver(\n",
        "  tf_env,\n",
        "  agent.collect_policy,\n",
        "  observers=replay_observer,\n",
        "  num_steps=collect_steps_per_iteration).run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "huGCDbO4GAF1"
      },
      "source": [
        "### トレーニングステップ用のデータ読み込み\n",
        "\n",
        "トラジェクトリ要素を再生バッファに追加した後は、再生バッファからトラジェクトリのバッチを読み取り、トレーニングステップの入力データとして使用できます。\n",
        "\n",
        "トレーニングループ内で再生バッファのトラジェクトリをトレーニングする方法の例を以下に示します。 "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gg8SUyXXnSMr"
      },
      "outputs": [],
      "source": [
        "# Read the replay buffer as a Dataset,\n",
        "# read batches of 4 elements, each with 2 timesteps:\n",
        "dataset = replay_buffer.as_dataset(\n",
        "    sample_batch_size=4,\n",
        "    num_steps=2)\n",
        "\n",
        "iterator = iter(dataset)\n",
        "\n",
        "num_train_steps = 10\n",
        "\n",
        "for _ in range(num_train_steps):\n",
        "  trajectories, _ = next(iterator)\n",
        "  loss = agent.train(experience=trajectories)\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "5_replay_buffers_tutorial.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
